import xarray as xr
import numpy as np
import pandas as pd

## ERAI land/sea mask
lsm = xr.open_dataset('MetData/ERA_mask.nc')
lsm = lsm['lsm'].values

## Data and file names
files = ['tp1_f32.nc','tp180_f32.nc','tmax_f32.nc','tmin_f32.nc','windspeed_f32.nc']
keys = ['tp','tp180','mx2t','mn2t','U10']
var = ['tp1','tp180','tmax','tmin','U10']
out = ['tp1_dist.pkl','tp180_dist.pkl','tmax_dist.pkl','tmin_dist.pkl','windspeed_dist.pkl']

## Save the meteorological distributions as vectors
for x in range(len(files)):
    data = xr.open_dataset(files[x])[keys[x]] # opens the data set
    data = data.where(lsm==1) # take only land and north of 60° S
    if keys[x] == 'tp' or keys[x] == 'tp180':
        data = (data.values*1000).flatten()
    elif keys[x] == 'mx2t' or keys[x] == 'mn2t':
        data = (data.values - 273.15).flatten()
    else:
        data = (data.values).flatten()
    data = data[~np.isnan(data)]
    globals()['%s' % var[x]] = pd.DataFrame(data={'%s'%(var[x]):data})
    pd.to_pickle(globals()['%s' % var[x]],out[x]) # save as data pickle
    del data,globals()['%s' % var[x]]